import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 加载数据
coordinates = np.load('output_coord.npy')  # 3D坐标 (N,3)
colors = np.load('output_color.npy')       # RGB颜色 (N,3) 范围0-255
segments = np.load('output_segment.npy')   # 分割标签 (N,)
print(segments)

# 检查数据形状
print(f"坐标形状: {coordinates.shape}")
print(f"颜色形状: {colors.shape}")
print(f"分割形状: {segments.shape}")

# 创建3D画布
fig = plt.figure(figsize=(18, 8))

# 原始RGB 3D图
ax1 = fig.add_subplot(121, projection='3d')
sc1 = ax1.scatter(
    coordinates[:, 0], coordinates[:, 1], coordinates[:, 2],
    c=colors/255.0,  # 归一化到0-1
    s=5,             # 点大小
    alpha=0.8,       # 透明度
    marker='o'       # 点形状
)
ax1.set_title('Original RGB (3D View)', fontsize=14)
ax1.set_xlabel('X axis')
ax1.set_ylabel('Y axis')
ax1.set_zlabel('Z axis')
ax1.view_init(elev=25, azim=45)  # 设置视角(仰角,方位角)

# 分割结果3D图
ax2 = fig.add_subplot(122, projection='3d')
unique_labels = np.unique(segments)
print(f"发现 {len(unique_labels)} 个分割类别")
print(set(unique_labels))
# 使用tab20 colormap (支持最多20个类别)
cmap = plt.get_cmap('tab20')
norm = plt.Normalize(vmin=min(unique_labels), vmax=max(unique_labels))

sc2 = ax2.scatter(
    coordinates[:, 0], coordinates[:, 1], coordinates[:, 2],
    c=segments,
    cmap=cmap,
    norm=norm,
    s=5,
    alpha=0.8,
    marker='o'
)
ax2.set_title('Segmentation Result (3D View)', fontsize=14)
ax2.set_xlabel('X axis')
ax2.set_ylabel('Y axis')
ax2.set_zlabel('Z axis')
ax2.view_init(elev=25, azim=45)

# 添加颜色条(仅对分割图)
cbar = fig.colorbar(sc2, ax=ax2, shrink=0.5)
cbar.set_label('Segment Labels')

plt.tight_layout()
plt.show()